import torch
import numpy as np
import cv2

def label_pad_norm(targets, pad, padded_size):
    '''
    :param targets: 标签
    :param pad: pad，思维向量，分别是
    :param padded_size:
    :return: [cx, cy, w, h]
    '''
    x1 = targets[:, 1]
    y1 = targets[:, 2]
    x2 = targets[:, 1] + targets[:, 3]
    y2 = targets[:, 2] + targets[:, 4]

    # Adjust for added padding
    x1 += pad[0]
    y1 += pad[2]
    x2 += pad[0]
    # x2 += pad[1]
    y2 += pad[2]
    # y2 += pad[3]

    labels = np.zeros((len(targets), 6), dtype=np.float32)
    labels[:, 1] = targets[:, 0]

    labels[:, 2] = ((x1 + x2) / 2) / padded_size[0]
    labels[:, 3] = ((y1 + y2) / 2) / padded_size[1]
    labels[:, 4] = targets[:, 3] / padded_size[0]
    labels[:, 5] = targets[:, 4] / padded_size[1]

    return labels

def decode_label_to_XYXY(labels, imageSize, outModel="XYXY"):
    '''
    解析yolo目标识别的标签，转换成xywh格式，并去归一化
    :param boxes:
    :param imageSize: 图像大小（w，h）
    '''
    boxes = labels[:, 1:]
    cx = boxes[:, 0] * imageSize[0]
    cy = boxes[:, 1] * imageSize[1]
    w = boxes[:, 2] * imageSize[0]
    h = boxes[:, 3] * imageSize[1]

    x1 = cx - w / 2
    y1 = cy - h / 2
    x2 = x1 + w
    y2 = y1 + h

    boxes_out = np.zeros((len(boxes), 5))
    boxes_out[:, 0] = labels[:, 0]

    boxes_out[:, 1] = x1
    boxes_out[:, 2] = y1
    boxes_out[:, 3] = x2
    boxes_out[:, 4] = y2

    return boxes_out.tolist()

def checklabel(labels, image, names):
    showimage = image * 255
    showimage = showimage.permute([1, 2, 0])
    showimage = np.asarray(showimage, dtype=np.uint8)

    c, h_img, w_img = image.shape

    labels = decode_label_to_XYXY(labels[:, 1:], (w_img, h_img))
    # labels = decode_label_to_XYXY(labels, (1, 1))

    for label, name in zip(labels, names):
        testimg = showimage.copy()
        label = label[1:]
        x1 = int(label[0])
        y1 = int(label[1])
        x2 = int(label[2])
        y2 = int(label[3])

        center_x = int((x1 + x2) / 2)
        center_y = int((y1 + y2) / 2)

        # testimg = cv2.line(testimg,(x1, y1), (x1, y2),(0,255,0),1)
        # testimg = cv2.line(testimg,(x1, y2), (x2, y2),(0,255,0),1)
        # testimg = cv2.line(testimg,(x2, y2), (x2, y1),(0,255,0),1)
        # testimg = cv2.line(testimg,(x2, y1), (x1, y1),(0,255,0),1)

        cv2.rectangle(testimg, (x1, y1), (x2, y2), (0,255,0))
        testimg = cv2.putText(testimg, name, (center_x, center_y), 1, 1.2, (0,255,0),)
        #

        testimg = cv2.resize(testimg, (800, 800))
        cv2.imshow("image", testimg)
        cv2.waitKey(0)